'''
    Author: Nebiyou Yismaw

    This is a python code that uses feature extraction techniques to perform
    panorama.

'''

import cv2
import numpy as np

def main():
    MAX_FEATURES = 500
    GOOD_MATCH_PERCENT = 0.15

    # Read reference image
    im1 = cv2.imread("scene1.jpg", cv2.IMREAD_COLOR)

    # Read image to be aligned
    im2 = cv2.imread("scene3.jpg", cv2.IMREAD_COLOR)

    # Convert images to grayscale
    im1Gray = cv2.cvtColor(im1, cv2.COLOR_BGR2GRAY)
    im2Gray = cv2.cvtColor(im2, cv2.COLOR_BGR2GRAY)

    # Detect ORB features and compute descriptors.
    orb = cv2.ORB_create(MAX_FEATURES)
    keypoints1, descriptors1 = orb.detectAndCompute(im1Gray, None)
    keypoints2, descriptors2 = orb.detectAndCompute(im2Gray, None)

    im1Keypoints = np.array([])
    im1Keypoints = cv2.drawKeypoints(im1, keypoints1, im1Keypoints, color=(0,0,255),flags=0)

    cv2.imshow("Keypoints",im1Keypoints)
    cv2.waitKey(0)

    # Match features.
    matcher = cv2.DescriptorMatcher_create(
                    cv2.DESCRIPTOR_MATCHER_BRUTEFORCE_HAMMING)
    matches = matcher.match(descriptors1, descriptors2, None)

    # Sort matches by score
    matches.sort(key=lambda x: x.distance, reverse=False)

    # Remove not so good matches
    numGoodMatches = int(len(matches) * GOOD_MATCH_PERCENT)
    matches = matches[:numGoodMatches]

    # Draw top matches
    imMatches = cv2.drawMatches(im1, keypoints1,
                                im2, keypoints2,
                                matches, None)

    cv2.imshow("Matchings obtained from the descriptor matcher",imMatches)
    cv2.waitKey(0)
    points1 = np.zeros((len(matches), 2), dtype=np.float32)
    points2 = np.zeros((len(matches), 2), dtype=np.float32)

    for i, match in enumerate(matches):
        points1[i, :] = keypoints1[match.queryIdx].pt
        points2[i, :] = keypoints2[match.trainIdx].pt

    # Find homography
    h, mask = cv2.findHomography(points2, points1, cv2.RANSAC)
    print("Homograhy matrix \n{}".format(h))

    # Use homography
    im1Height, im1Width, channels = im1.shape
    im2Height, im2Width, channels = im2.shape

    im2Aligned = cv2.warpPerspective(im2, h,
                                (im2Width + im1Width, im2Height))

    cv2.imshow("Second image aligned to first image obtained using homography and warping",im2Aligned)
    cv2.waitKey(0)
    # Stitch Image 1 with aligned image 2
    stitchedImage = np.copy(im2Aligned)
    stitchedImage[0:im1Height,0:im1Width] = im1

    cv2.imshow("Final Stitched Image",stitchedImage)
    cv2.waitKey(0)

    cv2.destroyAllWindows()

if __name__=="__main__":
    main()
